#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
快速 OCR 配置
提供多种 OCR 引擎选择，按速度排序
"""

import cv2
import numpy as np

class FastOCR:
    """快速 OCR 识别器，支持多种引擎"""
    
    def __init__(self, engine_priority=['tesseract', 'paddle', 'easyocr']):
        self.engine_priority = engine_priority
        self.engines = {}
        self.init_engines()
    
    def init_engines(self):
        """初始化可用的 OCR 引擎"""
        # 初始化 Tesseract
        try:
            import pytesseract
            self.engines['tesseract'] = pytesseract
            print("✓ Tesseract OCR 已加载")
        except ImportError:
            print("✗ Tesseract OCR 未安装")
        
        # 初始化 PaddleOCR
        try:
            from paddleocr import PaddleOCR
            self.paddle_ocr = PaddleOCR(use_angle_cls=False, lang='en', use_gpu=False)
            self.engines['paddle'] = self.paddle_ocr
            print("✓ PaddleOCR 已加载")
        except ImportError:
            print("✗ PaddleOCR 未安装")
        
        # 初始化 EasyOCR
        try:
            import easyocr
            self.easy_ocr_reader = easyocr.Reader(['en'])
            self.engines['easyocr'] = self.easy_ocr_reader
            print("✓ EasyOCR 已加载")
        except ImportError:
            print("✗ EasyOCR 未安装")
    
    def preprocess_image(self, roi):
        """预处理图像以提高识别效果"""
        # 转换为灰度图
        if len(roi.shape) == 3:
            gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        else:
            gray = roi
        
        # 应用高斯模糊去噪
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        
        # 自适应阈值处理
        thresh = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 11, 2
        )
        
        # 形态学操作清理噪点
        kernel = np.ones((2, 2), np.uint8)
        cleaned = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        
        return cleaned
    
    def recognize_with_tesseract(self, processed_image):
        """使用 Tesseract 识别"""
        try:
            # 尝试不同的 PSM 模式
            psm_modes = ['6', '7', '8', '13']
            for psm in psm_modes:
                text = self.engines['tesseract'].image_to_string(
                    processed_image, 
                    lang='eng', 
                    config=f'--psm {psm} --oem 3'
                )
                if text.strip():
                    return text.strip()
            return ""
        except Exception as e:
            print(f"Tesseract 识别错误: {e}")
            return ""
    
    def recognize_with_paddle(self, processed_image):
        """使用 PaddleOCR 识别"""
        try:
            results = self.engines['paddle'].ocr(processed_image, cls=False)
            if results and results[0]:
                texts = [line[1][0] for line in results[0] if line[1][1] > 0.3]
                return ' '.join(texts) if texts else ""
            return ""
        except Exception as e:
            print(f"PaddleOCR 识别错误: {e}")
            return ""
    
    def recognize_with_easyocr(self, processed_image):
        """使用 EasyOCR 识别"""
        try:
            results = self.engines['easyocr'].readtext(processed_image)
            texts = []
            for (bbox, text, prob) in results:
                if prob > 0.3:
                    texts.append(text.strip())
            return ' '.join(texts) if texts else ""
        except Exception as e:
            print(f"EasyOCR 识别错误: {e}")
            return ""
    
    def recognize(self, roi):
        """识别图像中的文本，按优先级尝试不同引擎"""
        if not self.engines:
            return "未安装 OCR 引擎"
        
        # 预处理图像
        processed_image = self.preprocess_image(roi)
        
        # 按优先级尝试不同引擎
        for engine_name in self.engine_priority:
            if engine_name in self.engines:
                if engine_name == 'tesseract':
                    result = self.recognize_with_tesseract(processed_image)
                elif engine_name == 'paddle':
                    result = self.recognize_with_paddle(processed_image)
                elif engine_name == 'easyocr':
                    result = self.recognize_with_easyocr(processed_image)
                
                if result:
                    return result
        
        return "未识别"
    
    def get_available_engines(self):
        """获取可用的 OCR 引擎列表"""
        return list(self.engines.keys())
    
    def benchmark_engines(self, test_image, iterations=10):
        """测试各引擎的速度"""
        import time
        
        results = {}
        for engine_name in self.engines:
            if engine_name == 'tesseract':
                start_time = time.time()
                for _ in range(iterations):
                    self.recognize_with_tesseract(test_image)
                elapsed = time.time() - start_time
                results[engine_name] = elapsed / iterations
            elif engine_name == 'paddle':
                start_time = time.time()
                for _ in range(iterations):
                    self.recognize_with_paddle(test_image)
                elapsed = time.time() - start_time
                results[engine_name] = elapsed / iterations
            elif engine_name == 'easyocr':
                start_time = time.time()
                for _ in range(iterations):
                    self.recognize_with_easyocr(test_image)
                elapsed = time.time() - start_time
                results[engine_name] = elapsed / iterations
        
        return results

# 使用示例
if __name__ == "__main__":
    # 创建快速 OCR 实例
    ocr = FastOCR()
    
    print(f"可用引擎: {ocr.get_available_engines()}")
    
    # 创建测试图像
    test_img = np.ones((100, 200), dtype=np.uint8) * 255
    cv2.putText(test_img, "TEST123", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, 0, 2)
    
    # 测试识别
    result = ocr.recognize(test_img)
    print(f"测试识别结果: {result}")
    
    # 性能测试
    if len(ocr.engines) > 1:
        print("\n性能测试:")
        benchmark_results = ocr.benchmark_engines(test_img, iterations=5)
        for engine, time_per_iter in benchmark_results.items():
            print(f"  {engine}: {time_per_iter:.4f} 秒/次")
